from datetime import datetime
from typing import Final, List
from enum import Enum, auto
from dateutil.relativedelta import relativedelta


class DatetimeUnits(Enum):
    DAY = auto()
    D = auto()
    WEEK = auto()
    W = auto
    MONTH = auto()
    M = auto()
    YEAR = auto()
    Y = auto()


class DatesRoundType(Enum):
    Before = auto()
    After = auto()


class TeaDatetime(object):
    _datetime_obj: datetime

    def __init__(
            self,
            date_string: str = '',
            date_format: str = '%Y-%m-%d',
            datetime_string: str = '',
            datetime_format: str = '%Y-%m-%d %H:%M:%S',
            datetime_obj: datetime = None
    ) -> "TeaDatetime":
        if datetime_obj:
            self._datetime_obj = datetime_obj
        elif date_string:
            self._datetime_obj = datetime.strptime(date_string, date_format)
        elif datetime_string:
            self._datetime_obj = datetime.strptime(datetime_string, datetime_format)
        else:
            self._datetime_obj = datetime.now()

    def _get_datetime_obj(self) -> 'datetime':
        return self._datetime_obj

    def change(self, offset: int, unit: DatetimeUnits) -> 'TeaDatetime':
        current_datetime_obj: datetime = self._get_datetime_obj()
        if unit in [DatetimeUnits.DAY, DatetimeUnits.D]:
            current_datetime_obj = current_datetime_obj + relativedelta(days=offset)
            pass
        elif unit in [DatetimeUnits.WEEK, DatetimeUnits.W]:
            current_datetime_obj = current_datetime_obj + relativedelta(days=offset * 7)
            pass
        elif unit in [DatetimeUnits.MONTH, DatetimeUnits.M]:
            current_datetime_obj = current_datetime_obj + relativedelta(months=offset)
            pass
        elif unit in [DatetimeUnits.YEAR, DatetimeUnits.Y]:
            current_datetime_obj = current_datetime_obj + relativedelta(years=offset)

        return TeaDatetime(datetime_obj=current_datetime_obj)

    def dates_round(self, length: int, unit: DatetimeUnits, dates_round_type: DatesRoundType) -> \
            List['TeaDatetime']:
        offset: int = 1 if dates_round_type == DatesRoundType.After else -1
        ret_list: list['TeaDatetime'] = []
        current_obj: TeaDatetime = self
        for i in range(length):
            current_obj: TeaDatetime = current_obj.change(offset * 1, unit)
            ret_list.append(current_obj)
        return ret_list

    def dates_round_str(self, length: int, unit: DatetimeUnits, dates_round_type: DatesRoundType, date_format: str) -> \
            List[str]:
        dates_list = self.dates_round(length, unit, dates_round_type)
        return [date_obj.strftime(date_format) for date_obj in dates_list]

    def strftime(self, fmt: str = '%Y-%m-%d %H:%M:%S') -> str:
        return self._get_datetime_obj().strftime(fmt)

    def __str__(self) -> str:
        return self._get_datetime_obj().__str__()


if __name__ == '__main__':
    print(TeaDatetime(date_string='2024-01-30').change(-10, DatetimeUnits.YEAR).change(1, DatetimeUnits.MONTH))
    print(TeaDatetime(date_string='2024-01-26').dates_round_str(10, DatetimeUnits.MONTH, DatesRoundType.After,
                                                                '%Y-%m-%d'))
